import pandas as pd
import json
import os
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("convert_university_ranking")

# 文件路径
current_dir = os.path.dirname(os.path.abspath(__file__))
csv_path = os.path.join(current_dir, '近8年软科中国大学排名.csv')
knowledge_path = os.path.join(current_dir, 'knowledge.json')

def load_existing_knowledge():
    """加载现有的知识库文件"""
    if os.path.exists(knowledge_path):
        try:
            with open(knowledge_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            logger.error("知识库文件格式错误，将创建新文件")
            return []
        except Exception as e:
            logger.error(f"读取知识库文件时出错: {e}")
            return []
    else:
        logger.info("知识库文件不存在，将创建新文件")
        return []

def convert_to_knowledge_entries():
    """将CSV数据转换为适合RAG的知识条目"""
    try:
        # 读取CSV文件
        logger.info(f"正在读取CSV文件: {csv_path}")
        df = pd.read_csv(csv_path)
        logger.info(f"CSV文件读取成功，共 {len(df)} 条记录")
        
        # 加载现有知识条目
        knowledge_data = load_existing_knowledge()
        initial_count = len(knowledge_data)
        logger.info(f"已加载现有知识库，包含 {initial_count} 条知识条目")
        
        # 获取最大ID
        max_id = 0
        if knowledge_data:
            try:
                max_id = max(item['id'] for item in knowledge_data)
            except KeyError:
                logger.warning("现有知识条目中没有id字段，将从1开始编号")
        
        # 按大学名称分组，整合多年数据
        universities = df['大学名称'].unique()
        logger.info(f"数据中包含 {len(universities)} 所不同的大学")
        
        # 创建知识条目
        new_entries = []
        for uni in universities:
            uni_data = df[df['大学名称'] == uni].sort_values('年份')
            
            # 基本信息（使用最新年份的数据）
            latest = uni_data.iloc[-1]
            uni_en = latest['大学英文名称']
            location = latest['所在地']
            uni_type = latest['类型']
            tags = latest['标签']
            latest_year = latest['年份']
            latest_rank = latest['排名']
            latest_score = latest['总分']
            
            # 创建基本介绍（包含常见搜索关键词）
            basic_info = (
                f"{uni}（{uni_en}）是一所位于{location}的{uni_type}类高校，"
                f"属于{tags}。"
            )
            
            # 添加最新排名信息（关键搜索信息优先展示）
            latest_info = (
                f"在{latest_year}年软科中国大学排名中，{uni}位列第{latest_rank}名，总分为{latest_score}分。"
            )
            
            # 添加历年排名趋势信息
            rankings = []
            years = []
            for _, row in uni_data.iterrows():
                year = row['年份']
                years.append(year)
                rank = row['排名']
                score = row['总分']
                rankings.append(f"{year}年排名第{rank}位（总分{score}分）")
            
            # 分析排名趋势
            first_rank = uni_data.iloc[0]['排名']
            last_rank = uni_data.iloc[-1]['排名']
            trend_desc = ""
            if first_rank > last_rank:
                trend_desc = f"从{years[0]}年的第{first_rank}名上升到{years[-1]}年的第{last_rank}名，"
                trend_desc += "排名整体呈上升趋势。"
            elif first_rank < last_rank:
                trend_desc = f"从{years[0]}年的第{first_rank}名下降到{years[-1]}年的第{last_rank}名，"
                trend_desc += "排名整体呈下降趋势。"
            else:
                trend_desc = f"在{years[0]}至{years[-1]}年期间，排名保持稳定在第{first_rank}名。"
            
            trend_info = f"该校在{years[0]}至{years[-1]}年间的排名变化：{trend_desc}"
            
            # 详细历年排名记录
            ranking_detail = "历年排名详情：" + "；".join(rankings) + "。"
            
            # 为不同查询意图添加结构化关键词段落
            keywords_section = (
                f"关键词：大学排名 {uni} {uni_en} {location} {uni_type} {tags} "
                f"软科排名 中国大学排名 {latest_year}年 第{latest_rank}名 {latest_score}分"
            )
            
            # 组合完整的知识条目
            knowledge_text = (
                f"{basic_info}\n{latest_info}\n{trend_info}\n{ranking_detail}\n{keywords_section}"
            )
            
            # 添加到新条目列表
            max_id += 1
            new_entries.append({
                "id": max_id,
                "text": knowledge_text
            })
            
            # 为每个年份创建单独的排名记录，方便直接检索特定年份的排名
            for _, row in uni_data.iterrows():
                year = row['年份']
                rank = row['排名']
                score = row['总分']
                
                year_specific_text = (
                    f"{year}年软科中国大学排名：{uni}（{uni_en}）位列第{rank}名，"
                    f"总分为{score}分。该校是一所位于{location}的{uni_type}类高校，属于{tags}。"
                )
                
                max_id += 1
                new_entries.append({
                    "id": max_id,
                    "text": year_specific_text
                })
        
        # 创建年度排名总结知识条目
        for year in df['年份'].unique():
            year_data = df[df['年份'] == year].sort_values('排名').head(10)
            
            # 生成前10名名单
            top10_list = []
            for _, row in year_data.iterrows():
                uni = row['大学名称']
                rank = row['排名']
                score = row['总分']
                top10_list.append(f"第{rank}名：{uni}（{score}分）")
            
            year_summary = (
                f"{year}年软科中国大学排名前10名院校：\n" + 
                "；".join(top10_list) + "。\n" +
                f"关键词：{year}年 软科排名 中国大学排名 前10名 前十名 排行榜"
            )
            
            max_id += 1
            new_entries.append({
                "id": max_id,
                "text": year_summary
            })
        
        # 添加到知识库
        knowledge_data.extend(new_entries)
        logger.info(f"已创建 {len(new_entries)} 条新知识条目")
        
        # 保存更新后的知识库
        with open(knowledge_path, 'w', encoding='utf-8') as f:
            json.dump(knowledge_data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"知识库已更新，当前共有 {len(knowledge_data)} 条知识条目")
        logger.info(f"新增 {len(knowledge_data) - initial_count} 条知识条目")
        
        # 打印一条示例条目
        if new_entries:
            logger.info(f"示例条目: {new_entries[0]['text'][:200]}...")
        
        return True
    
    except Exception as e:
        logger.error(f"转换过程中出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return False

if __name__ == "__main__":
    logger.info("开始转换大学排名数据到知识库")
    success = convert_to_knowledge_entries()
    if success:
        logger.info("转换完成")
    else:
        logger.error("转换失败") 